"""缓存相关封装"""

import pickle
import tempfile
from os.path import join

from diskcache import Cache

from lib.pressure_task import get_pressure_task_sig
from lib.global_share import GlobalShare, note_taker, error_note_taker


cache = Cache(GlobalShare.CACHE_DIR)


def set_task_id_set(data_set: set):
    """设置 task_id_set 缓存"""

    cache.set("task_id_set", data_set)


def get_task_id_set() -> set:
    """获取 task_id_set 缓存"""

    return cache.get("task_id_set", default=set())


def set_locust_processes_cache(value):
    """设置 locust 进程信息缓存"""

    cache.set("locust_processes", value)


def get_cached_locust_processes():
    """获取 locust 进程信息缓存"""

    return cache.get("locust_processes", {})


def del_all_cacahed_locust_processes():
    """删除缓存的 locust 进程信息缓存"""

    note_taker.info("[del_all_cacahed_locust_processes] start to delete")
    cache.delete("locust_processes")


def del_cached_locust_processes(pressure_task_id):
    """删除缓存的 locust 进程信息"""

    pressure_task_sig = get_pressure_task_sig(pressure_task_id=pressure_task_id)
    locust_processes: dict = cache.get("locust_processes", {})
    if pressure_task_sig in locust_processes:
        locust_processes.pop(pressure_task_sig)
        cache.set("locust_processes", locust_processes)


def update_cached_locust_processes(pressure_task_id, node_name, node_info):
    """更新缓存的 locust 进程信息"""

    if node_name not in ["master", "workers"]:
        error_note_taker.error(
            f"[update_cached_locust_processes] node name error: {node_name}"
        )
        return
    pressure_task_sig = get_pressure_task_sig(pressure_task_id=pressure_task_id)
    locust_processes: dict = cache.get("locust_processes", {})
    if pressure_task_sig in locust_processes:
        locust_processes[pressure_task_sig][node_name] = node_info
        cache.set("locust_processes", locust_processes)


def get_pressure_task_cached_info(pressure_task_id):
    """获取缓存的压测任务信息"""

    pressure_task_sig = get_pressure_task_sig(pressure_task_id=pressure_task_id)
    locust_processes = get_cached_locust_processes()
    if pressure_task_sig not in locust_processes:
        return {"master": {}, "workers": []}
    return locust_processes[pressure_task_sig]


def cache_master_info(pressure_task_id: str, info: dict):
    """缓存 master 信息"""

    note_taker.info(f"[cache_master_info] cache master info: {info}")
    pressure_task_sig = get_pressure_task_sig(pressure_task_id=pressure_task_id)
    _info = get_pressure_task_cached_info(pressure_task_id)
    _info["master"] = info
    locust_processes = get_cached_locust_processes()
    locust_processes[pressure_task_sig] = _info
    set_locust_processes_cache(locust_processes)
    note_taker.info(f"[cache_master_info] master info cached: {locust_processes}")


def cache_worker_info(pressure_task_id: str, info: dict):
    """缓存 worker 信息"""

    note_taker.info(f"[cache_worker_info] info: {info}")
    pressure_task_sig = get_pressure_task_sig(pressure_task_id=pressure_task_id)
    _info = get_pressure_task_cached_info(pressure_task_id)
    _info["workers"].append(info)
    locust_processes = get_cached_locust_processes()
    locust_processes[pressure_task_sig] = _info
    set_locust_processes_cache(locust_processes)
    note_taker.info(f"[cache_worker_info] worker info cached: {locust_processes}")


def cache_python_interpreter_path(project_name, path):
    """缓存项目的 python 解释器路径"""

    cache_dir = tempfile.gettempdir()
    cache_file = join(cache_dir, f"{project_name}_interpreter_cache.pkl")
    note_taker.info(f"[cache_python_interpreter_path] cache_file: {cache_file}")
    with open(cache_file, "wb") as file:
        pickle.dump(path, file)


def load_python_interpreter_path(project_name):
    """加载项目缓存的 python 解释器路径"""

    cache_dir = tempfile.gettempdir()
    cache_file = join(cache_dir, f"{project_name}_interpreter_cache.pkl")
    note_taker.info(f"[load_python_interpreter_path] cache_file: {cache_file}")
    try:
        with open(cache_file, "rb") as file:
            path = pickle.load(file)
            return path
    except FileNotFoundError:
        return None
